'''
1 获取所有图片的embedding
2 计算embedding最相似的图片，余弦相似度，越接近1越相似
'''
import os
# from os.path import join as osj
from util_for_os import ose,osj
os.environ['CUDA_VISIBLE_DEVICES'] = '3'

import faiss,pdb,torch

from transformers import AutoModel, AutoProcessor
from transformers.image_utils import load_image
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd
import shutil
from tqdm import tqdm
import numpy as np

from demo_dinov3 import get_dino3_processor,get_pooled_feature_by_dino3
from util_flux import horizontal_concat_images


SIGLIP_PATH = '/data/models/siglip2-so400m-patch16-512'

def get_siglip2():
    processor = AutoProcessor.from_pretrained(SIGLIP_PATH,)
    model = AutoModel.from_pretrained(SIGLIP_PATH, device_map="cuda").eval()
    return processor,model.cuda()
 
def extract_feature_dino( filepath='',file_pil=None , image_processor=None, model=None ):
    # prepare input image
    image = load_image(filepath) if file_pil is None else file_pil
    inputs = image_processor(images=[image], 
                             return_tensors="pt").to(model.device)

    with torch.no_grad():
        image_embeddings = model.get_image_features(**inputs)  # 1 1152
    
    return image_embeddings.cpu().numpy()

date = '0908'

categories = ['coat']

for category in categories:
    
    img_dir = f'/mnt/nas/datasets/diction/{category}{date}_img_clo'
    assert ose(img_dir)
    assert ose( osj(img_dir , 'names.txt') )


# dino_processor,dino3 = get_dino3_processor()
# sig_processor, siglip = get_siglip2()

save_dir_dino = './compare_dino'
save_dir_sig = './compare_sig'
os.makedirs(save_dir_dino,exist_ok=True)
os.makedirs(save_dir_sig,exist_ok=True)
for category in categories:
    
    img_dir = f'/mnt/nas/datasets/diction/{category}{date}_img_clo'
    assert ose(img_dir)
    assert ose( osj(img_dir , 'names.txt') )

    with open( osj(img_dir , 'names.txt')) as f:
        filenames = f.readlines()
        filenames = [f.strip() for f in filenames if f.endswith('.jpg\n') ]
    # pdb.set_trace()
    tmp_save_emb_pt_sig = os.path.join( img_dir , 'tmp_save_emb_pt.pt' )
    tmp_save_emb_pt = os.path.join( img_dir , 'tmp_save_emb_pt_dino3.pt' )

    # Step 1: Extract embeddings for all images and save to CSV
    embeddings_sig = []
    abs_paths_sig = []    
    embeddings = []
    abs_paths = []

    # Check if cached embeddings exist
    assert ose(tmp_save_emb_pt_sig)
    assert ose(tmp_save_emb_pt)
    
    print("Loading cached embeddings...")
    saved_emb_sig = torch.load(tmp_save_emb_pt_sig)
    embeddings_sig = saved_emb_sig['embeddings'].numpy()
    abs_paths_sig = saved_emb_sig['filepath']        
    
    saved_emb = torch.load(tmp_save_emb_pt)
    embeddings = saved_emb['embeddings'].numpy()
    abs_paths = saved_emb['filepath']

        
    # Step 2: Calculate pairwise similarities and move similar/different images
    threshold = 0.98  # Similarity threshold
    sim_paths = set()

    for i in tqdm(range(len(embeddings))):
        for j in range(i + 1, min(i+50 , len(embeddings))):
            sim_sig = cosine_similarity(embeddings_sig[i], embeddings_sig[j])[0][0]
            sim = cosine_similarity(embeddings[i], embeddings[j])[0][0]
            
            check_sig = sim_sig > threshold
            check_dino = sim > threshold
        
            if check_sig or check_dino :
                test_img = horizontal_concat_images( [ load_image( abs_paths[i] ) , 
                                                        load_image( abs_paths[j] ) ] )

                save_name = f'sig_{check_sig}_{sim_sig:.2f}_dino_{check_dino}_{sim:.2f}'          
                concat_save_name = f'{save_name}_{os.path.basename(abs_paths[i])}__{os.path.basename(abs_paths[j])}.jpg'

                if check_sig and not check_dino:
                    print('sig',check_sig,'dino',check_dino)  
                    # 保存到save_dir_sig
                    concat_save_path = os.path.join(save_dir_sig, concat_save_name)
                    test_img.save(concat_save_path)
                if check_dino and not check_sig:
                    print('sig',check_sig,'dino',check_dino)
                    # 保存到save_dir_dino
                    concat_save_path = os.path.join(save_dir_dino, concat_save_name)
                    test_img.save(concat_save_path)

            # if sim > threshold or sim_sig > threshold:
            #     from util_flux import horizontal_concat_images
            #     horizontal_concat_images( [ load_image( abs_paths[i] ) , 
            #                             load_image( abs_paths[j] ) ] ).save( 'tmp_dino3.jpg' )
            #     sim_paths.add( abs_paths[j] )


